package ffte.test.types

import spinal.core._
import spinal.core.sim._
import ffte.types._

/**
 * VecComplex正确测试套件
 * 使用正确的字段访问方式：通过.d字段访问Vec[FixComplex]
 */
class VecComplexBasicCorrectTest extends Component {
  val io = new Bundle {
    val real_in = in Vec(SInt(16 bits), 4)  // 4个复数的实部
    val imag_in = in Vec(SInt(16 bits), 4)  // 4个复数的虚部
    val real_out = out Vec(SInt(16 bits), 4) // 输出实部
    val imag_out = out Vec(SInt(16 bits), 4) // 输出虚部
    val select = in UInt(2 bits)  // 选择要输出的复数
  }

  val testArea = new Area {
    // 创建VecComplex实例
    val vecComplex = VecComplex(4, 16, -8)  // 4个元素，16位，8位小数
    val result = VecComplex(4, 16, -8)
    
    // 连接输入到VecComplex - 通过.d字段访问
    for (i <- 0 until 4) {
      vecComplex.d(i).re.d.raw := io.real_in(i)
      vecComplex.d(i).im.d.raw := io.imag_in(i)
    }
    
    // 简单的元素操作：每个元素乘以2
    for (i <- 0 until 4) {
      val doubled = vecComplex.d(i) + vecComplex.d(i)  // 复数自加等于乘以2
      result.d(i).re.d.raw := doubled.re.d.raw
      result.d(i).im.d.raw := doubled.im.d.raw
    }
    
    // 输出所有元素
    for (i <- 0 until 4) {
      io.real_out(i) := result.d(i).re.d.raw
      io.imag_out(i) := result.d(i).im.d.raw
    }
  }
}

/**
 * VecComplex元素访问和操作测试
 */
class VecComplexElementCorrectTest extends Component {
  val io = new Bundle {
    val real_in = in Vec(SInt(16 bits), 3)
    val imag_in = in Vec(SInt(16 bits), 3)
    val real_out1 = out SInt(16 bits)
    val imag_out1 = out SInt(16 bits)
    val real_out2 = out SInt(16 bits)
    val imag_out2 = out SInt(16 bits)
    val index = in UInt(2 bits)
  }

  val elementArea = new Area {
    // 创建VecComplex
    val vec = VecComplex(3, 16, -8)
    val element1 = FixComplex(16, -8)
    val element2 = FixComplex(16, -8)
    
    // 填充VecComplex
    for (i <- 0 until 3) {
      vec.d(i).re.d.raw := io.real_in(i)
      vec.d(i).im.d.raw := io.imag_in(i)
    }
    
    // 元素访问和操作 - 通过.d字段
    element1 := vec.d(io.index)  // 动态索引访问
    element2 := vec.d(1)  // 静态索引访问
    
    // 元素操作：访问的元素乘以j（旋转90度）
    val rotated1 = element1.j
    val rotated2 = element2.j
    
    // 连接输出
    io.real_out1 := rotated1.re.d.raw
    io.imag_out1 := rotated1.im.d.raw
    io.real_out2 := rotated2.re.d.raw
    io.imag_out2 := rotated2.im.d.raw
  }
}

/**
 * VecComplex与FixComplex交互测试
 */
class VecComplexInteractionCorrectTest extends Component {
  val io = new Bundle {
    val real_scalar = in SInt(16 bits)
    val imag_scalar = in SInt(16 bits)
    val real_vec = in Vec(SInt(16 bits), 2)
    val imag_vec = in Vec(SInt(16 bits), 2)
    val real_out = out Vec(SInt(32 bits), 2)
    val imag_out = out Vec(SInt(32 bits), 2)
  }

  val interactionArea = new Area {
    // 创建标量复数
    val scalar = FixComplex(16, -8)
    scalar.re.d.raw := io.real_scalar
    scalar.im.d.raw := io.imag_scalar
    
    // 创建向量复数
    val vecComplex = VecComplex(2, 16, -8)
    for (i <- 0 until 2) {
      vecComplex.d(i).re.d.raw := io.real_vec(i)
      vecComplex.d(i).im.d.raw := io.imag_vec(i)
    }
    
    // 标量与向量每个元素相乘
    for (i <- 0 until 2) {
      val product = scalar * vecComplex.d(i)  // 复数乘法
      io.real_out(i) := product.re.d.raw
      io.imag_out(i) := product.im.d.raw
    }
  }
}

/**
 * VecComplex特殊操作测试
 */
class VecComplexSpecialOpsTest extends Component {
  val io = new Bundle {
    val real_in = in Vec(SInt(16 bits), 4)
    val imag_in = in Vec(SInt(16 bits), 4)
    val real_out = out Vec(SInt(16 bits), 4)
    val imag_out = out Vec(SInt(16 bits), 4)
  }

  val specialArea = new Area {
    // 创建VecComplex
    val vec = VecComplex(4, 16, -8)
    for (i <- 0 until 4) {
      vec.d(i).re.d.raw := io.real_in(i)
      vec.d(i).im.d.raw := io.imag_in(i)
    }
    
    // 特殊操作：共轭所有元素
    val conjugated = vec.conj
    
    // 输出共轭结果
    for (i <- 0 until 4) {
      io.real_out(i) := conjugated.d(i).re.d.raw
      io.imag_out(i) := conjugated.d(i).im.d.raw
    }
  }
}

/**
 * 完整测试套件仿真
 */
object VecComplexCorrectTestSim {
  def main(args: Array[String]): Unit = {
    // 测试1：基本VecComplex操作
    println("开始VecComplex基本测试...")
    SimConfig.withWave.compile(new VecComplexBasicCorrectTest()).doSim { dut =>
      dut.clockDomain.forkStimulus(period = 10)
      
      // 测试数据：4个复数
      val testReals = Array(1.0, 2.0, 3.0, 4.0)
      val testImags = Array(1.0, 2.0, 3.0, 4.0)
      
      // 填充输入
      for (i <- 0 until 4) {
        dut.io.real_in(i) #= Math.round(testReals(i) * 128)
        dut.io.imag_in(i) #= Math.round(testImags(i) * 128)
      }
      
      // 选择输出第一个元素
      dut.io.select #= 0
      dut.clockDomain.waitSampling(1)
      
      val out_real = dut.io.real_out(0).toInt
      val out_imag = dut.io.imag_out(0).toInt
      println(f"VecComplex[0]乘以2: ($out_real/128 + ${out_imag}/128i)")
      println(f"期望: (2+2i)")
      
      // 测试选择第三个元素
      dut.io.select #= 2
      dut.clockDomain.waitSampling(1)
      
      val out_real3 = dut.io.real_out(0).toInt
      val out_imag3 = dut.io.imag_out(0).toInt
      println(f"VecComplex[2]乘以2: ($out_real3/128 + ${out_imag3}/128i)")
      println(f"期望: (6+6i)")
    }
    
    // 测试2：元素访问和操作
    println("\n开始VecComplex元素访问测试...")
    SimConfig.withWave.compile(new VecComplexElementCorrectTest()).doSim { dut =>
      dut.clockDomain.forkStimulus(period = 10)
      
      // 测试数据：3个复数
      val testReals = Array(1.0, 2.0, 3.0)
      val testImags = Array(4.0, 5.0, 6.0)
      
      for (i <- 0 until 3) {
        dut.io.real_in(i) #= Math.round(testReals(i) * 128)
        dut.io.imag_in(i) #= Math.round(testImags(i) * 128)
      }
      
      // 测试动态索引访问（索引1）
      dut.io.index #= 1
      dut.clockDomain.waitSampling(1)
      
      val real1 = dut.io.real_out1.toInt
      val imag1 = dut.io.imag_out1.toInt
      println(f"VecComplex[1]旋转90度: ($real1/128 + ${imag1}/128i)")
      println(f"原始: (2+5i), 旋转后期望: (-5+2i)")
      
      // 测试静态索引访问（索引1）
      val real2 = dut.io.real_out2.toInt
      val imag2 = dut.io.imag_out2.toInt
      println(f"VecComplex[1]静态访问旋转: ($real2/128 + ${imag2}/128i)")
    }
    
    // 测试3：VecComplex与FixComplex交互
    println("\n开始VecComplex与FixComplex交互测试...")
    SimConfig.withWave.compile(new VecComplexInteractionCorrectTest()).doSim { dut =>
      dut.clockDomain.forkStimulus(period = 10)
      
      // 标量复数：2+3i
      dut.io.real_scalar #= Math.round(2.0 * 128)
      dut.io.imag_scalar #= Math.round(3.0 * 128)
      
      // 向量复数：[1+2i, 3+4i]
      dut.io.real_vec(0) #= Math.round(1.0 * 128)
      dut.io.imag_vec(0) #= Math.round(2.0 * 128)
      dut.io.real_vec(1) #= Math.round(3.0 * 128)
      dut.io.imag_vec(1) #= Math.round(4.0 * 128)
      
      dut.clockDomain.waitSampling(1)
      
      val real0 = dut.io.real_out(0).toInt
      val imag0 = dut.io.imag_out(0).toInt
      val real1 = dut.io.real_out(1).toInt
      val imag1 = dut.io.imag_out(1).toInt
      
      println(f"标量(2+3i) × VecComplex[0](1+2i) = ($real0/${128*128} + ${imag0}/${128*128}i)")
      println(f"标量(2+3i) × VecComplex[1](3+4i) = ($real1/${128*128} + ${imag1}/${128*128}i)")
      
      // 手工计算验证：
      // (2+3i)(1+2i) = (2-6) + (4+3)i = -4 + 7i
      // (2+3i)(3+4i) = (6-12) + (8+9)i = -6 + 17i
      println(f"期望结果: (-4+7i) 和 (-6+17i)")
      println(f"实际值: (${real0/(128.0*128.0)}%+6.3f + ${imag0/(128.0*128.0)}%+6.3fi) 和 (${real1/(128.0*128.0)}%+6.3f + ${imag1/(128.0*128.0)}%+6.3fi)")
    }
    
    // 测试4：VecComplex特殊操作
    println("\n开始VecComplex特殊操作测试...")
    SimConfig.withWave.compile(new VecComplexSpecialOpsTest()).doSim { dut =>
      dut.clockDomain.forkStimulus(period = 10)
      
      // 测试数据：4个复数
      val testReals = Array(1.0, 2.0, 3.0, 4.0)
      val testImags = Array(1.0, -2.0, 3.0, -4.0)
      
      for (i <- 0 until 4) {
        dut.io.real_in(i) #= Math.round(testReals(i) * 128)
        dut.io.imag_in(i) #= Math.round(testImags(i) * 128)
      }
      
      dut.clockDomain.waitSampling(1)
      
      println("VecComplex共轭操作结果：")
      for (i <- 0 until 4) {
        val real_conj = dut.io.real_out(i).toInt
        val imag_conj = dut.io.imag_out(i).toInt
        println(f"  元素$i: ($real_conj/128 + ${imag_conj}/128i) [期望共轭: (${testReals(i)}%+3.1f${-testImags(i)}%+3.1fi)]")
      }
    }
    
    println("\nVecComplex完整测试完成！")
    println("关键观察：")
    println("1. VecComplex包含.d字段，是FixComplex的Vec")
    println("2. 通过vec.d(i)访问各个元素")
    println("3. 支持元素级运算和特殊操作")
    println("4. 可以与标量FixComplex进行交互")
    println("5. 支持各种复数操作（共轭、旋转、交换等）")
  }
}